from pickletools import optimize

import torch
import torchvision.datasets as dataset
import torchvision.transforms as transforms
import torch.utils.data as data_utils
from openpyxl.styles.builtins import output

from CNN import CNN

# 数据加载
train_data = dataset.MNIST(
    root = "mnist",
    # 加载训练数据
    train = True,
    # 转成张量
    transform = transforms.ToTensor(),
    download = True
)

test_data = dataset.MNIST(
    root = "mnist",
    # 加载测试集
    train = False,
    transform = transforms.ToTensor(),
    download = True
)

# 分批加载
train_loader = data_utils.DataLoader(
    dataset= train_data,
    batch_size=64,
    shuffle=True)

test_loader = data_utils.DataLoader(
    dataset= test_data,
    batch_size=64,
    shuffle=True)

cnn = CNN()

# 损失函数
loss_func = torch.nn.CrossEntropyLoss()

# 优化函数
optimizer = torch.optim.Adam(cnn.parameters(),lr=0.01)

# 训练过程
for epoch in range(3):
    for index,(images,labels) in enumerate(train_loader):
        # print(index)
        # print(images)
        # print(labels)
        # 前向传播
        outputs = cnn(images)

        loss = loss_func(outputs,labels)

        # 清空梯度
        optimizer.zero_grad()
        # 反向传播
        loss.backward()
        optimizer.step()
        print("当前epoch为{}，当前批次为{}/{}，loss为{}".
              format(epoch+1,index+1,len(train_data)//64,loss.item()))

    # 测试集验证
    loss_test = 0
    rightValue = 0
    for index2,(images,labels) in enumerate(test_loader):
        outputs = cnn(images)
        # print(outputs)
        # print(outputs.size())
        # print(labels)
        # print(labels.size())
        loss_test += loss_func(outputs,labels)
        _,pred = outputs.max(1)
        # print(pred)
        rightValue += (pred==labels).sum().item()
        print("测试集：当前epoch为{}，当前批次为{}/{}，loss为{}，准确率为{}".
              format(epoch + 1,index2 + 1,len(test_data)//64,loss_test,rightValue/len(test_data)))

torch.save(cnn,"model/mnist_model.pkl")
